import os
from typing import Dict, Tuple
from langchain_core.messages import SystemMessage, HumanMessage
from config.crypto_config import DOC_PROMPT_TEMPLATES, DOC_KEY_EXTRACTION, SYSTEM_MESSAGE_TEMPLATES
from services.file.llm_service import LLMService

class DocExtractionService:
    """文档提取服务类，处理密码应用方案文档的信息提取"""
    
    def __init__(self):
        # 初始化LLM服务，使用更好的错误处理
        try:
            self.llm_service = LLMService()
            print("DocExtractionService: LLM服务初始化成功")
        except Exception as e:
            print(f"DocExtractionService: LLM服务初始化失败: {e}")
            self.llm_service = None
        
        self.doc_prompts = DOC_PROMPT_TEMPLATES
        self.doc_key_extraction = DOC_KEY_EXTRACTION
    
    def extract_information_from_doc(self, doc_path: str) -> Dict[str, str]:
        """
        从文档中提取信息
        
        Args:
            doc_path (str): 文档路径
            
        Returns:
            Dict[str, str]: 提取的信息字典
        """
        # 这里需要实现文档加载和解析逻辑
        # 由于原项目使用了自定义的RapidOCRDocLoader，这里提供一个简化版本
        # 您可以根据需要移植完整的文档加载逻辑
        
        # 模拟文档内容提取
        doc_content = self._load_document(doc_path)
        return self._extract_structured_info(doc_content)
    
    def _load_document(self, doc_path: str) -> str:
        """
        加载文档内容
        
        Args:
            doc_path (str): 文档路径
            
        Returns:
            str: 文档内容
        """
        # 这里应该实现实际的文档加载逻辑
        # 可以移植原项目的RapidOCRDocLoader
        try:
            with open(doc_path, 'r', encoding='utf-8') as f:
                return f.read()
        except Exception as e:
            raise Exception(f"文档加载失败: {str(e)}")
    
    def _extract_structured_info(self, doc_content: str) -> Dict[str, str]:
        """
        从文档内容中提取结构化信息
        
        Args:
            doc_content (str): 文档内容
            
        Returns:
            Dict[str, str]: 结构化信息
        """
        crypto_eval_information = {
            "KEY": "",
            "AUTH": "",
            "ACCESS_CONTROL": "",
            "KEY_DATA_TRANSPORT_CONFIDENTIALITY_AND_INTEGRITY": "",
            "KEY_DATA_STORE_CONFIDENTIALITY": "",
            "KEY_DATA_STORE_INTEGRITY": "",
            "KEY_DATA_SAFE_MARK_INTEGRITY": "",
            "NON-REPUDIATION": "",
        }
        
        # 这里应该实现基于章节编号的信息提取逻辑
        # 可以移植原项目的章节识别和内容提取逻辑
        
        return crypto_eval_information
    
    def process_document_with_llm(self, doc_path: str) -> Tuple[Dict[str, str], Dict[str, str]]:
        """
        使用LLM处理文档并提取信息
        
        Args:
            doc_path (str): 文档路径
            
        Returns:
            Tuple[Dict[str, str], Dict[str, str]]: (原始信息, 结构化信息)
        """
        # 获取原始文档信息
        crypto_eval_information = self.extract_information_from_doc(doc_path)
        
        # 使用LLM进行结构化处理
        crypto_eval_extract = {
            "KEY": "",
            "AUTH": "",
            "ACCESS_CONTROL": "",
            "KEY_DATA_TRANSPORT_CONFIDENTIALITY_AND_INTEGRITY": "",
            "KEY_DATA_STORE_CONFIDENTIALITY": "",
            "KEY_DATA_STORE_INTEGRITY": "",
            "KEY_DATA_SAFE_MARK_INTEGRITY": "",
            "NON-REPUDIATION": "",
        }
        
        if crypto_eval_information and self.llm_service:
            for key in crypto_eval_extract:
                if crypto_eval_information[key]:
                    prompt_key = f"{key}_INFORMATION"
                    if prompt_key in self.doc_prompts:
                        prompt_template = self.doc_prompts[prompt_key]
                        prompt = prompt_template.format(doc_paragraph=crypto_eval_information[key])
                        
                        # 使用LLM处理，添加错误处理
                        try:
                            llm_response = self.llm_service.generate_response(prompt)
                            crypto_eval_extract[key] = llm_response
                        except Exception as e:
                            print(f"LLM处理 {key} 失败: {e}")
                            crypto_eval_extract[key] = f"处理失败: {str(e)}"
        
        return crypto_eval_information, crypto_eval_extract 